
# Set up the packages, options, and working directory path
pacman::p_load(tidyverse, RColorBrewer)
set.seed(42)
options(scipen=999)
if (Sys.info()[['user']] == "alal") {
  root = "/home/alal/Dropbox/1_Research/ElecAdminFunding/ElecAdminFundingReplication"
  theme_set(lal_plot_theme())
} else {
  root = "~/Dropbox/ElecAdminFunding/ElecAdminFundingReplication"
  theme_set(theme_minimal())
}

blues = brewer.pal(8, 'Blues')[c(5, 8)]

# Set up the SDID output for plotting
load(file = file.path(root, "tmp/allsynth.RData"), verbose = TRUE)

# Set up a function that extracts a counterfactual from our output
get_counterfactual = function(est) {
  setup = attr(est, 'setup')
  weights = attr(est, 'weights')
  Y = setup$Y
  N0 = setup$N0; N1 = nrow(Y) - N0
  T0 = setup$T0; T1 = ncol(Y) - T0
  lambda.synth = c(weights$lambda, rep(0, T1))
  lambda.target = c(rep(0, T0), rep(1 / T1, T1))
  omega.synth = c(weights$omega, rep(0, N1))
  omega.target = c(rep(0, N0), rep(1 / N1, N1))
  intercept.offset = c((omega.target - omega.synth) %*% Y %*% lambda.synth)
  obs.trajectory = as.numeric(omega.target %*% Y)
  syn.trajectory = as.numeric(omega.synth %*% Y) + intercept.offset
  return(list(syn=syn.trajectory, obs=obs.trajectory))
}


##
# Plots Comparing Treated County vs SDID Counterfactual
##

# Set up a function that plots treated county against the SDID counterfactual
make_levels_plot = function(est_list, y_axis_title="XX"){
  data.frame(year = seq(1992, 2020, by=4),
    y_1 = get_counterfactual(est_list$ests[["Synthetic DID (SDID)"]])$obs,
    y_0 = get_counterfactual(est_list$ests[["Synthetic DID (SDID)"]])$syn) %>%
    pivot_longer(
      cols = starts_with("y_"),
      names_to = "treat",
      values_to = "y"
    ) %>%
    mutate(treat = case_when(
        treat=="y_1" ~ "Grant Recipient",
        treat=="y_0" ~ "Synthetic Grant Recipient")) %>%
    ggplot(aes(x=year, y=y, color=treat, linetype=treat)) +
      geom_vline(xintercept = 2018, linetype = 'dotted', linewidth=1) +
      geom_line(linewidth=1.5) + geom_point(size=2) +
      ylab(y_axis_title) + xlab("Year") + #ylim(-1, 1) +
      scale_x_continuous(limits = c(1991, 2021),
        breaks = seq(1992, 2020, by=4)) +
      theme(text = element_text(size=20, face="bold", color="black"),
        legend.position="bottom", legend.title= element_blank(),
        panel.grid.minor = element_blank()) +
      theme(panel.grid.minor = element_blank()) +
      scale_color_manual(values = rev(blues), name = "")  +
      guides(linetype = "none") %>%
    return()
}

# Make the simple vote share SDID plot
fig_vs = make_levels_plot(allsynth_dem, y_axis_title = "Dem Vote Share (%)")
ggsave(file.path(root, "output/sdid_levels_demvs_figure.pdf"),
  plot = fig_vs, height = 8, width = 10)

# Make the simple turnout SDID plot
fig_turnout = make_levels_plot(allsynth_tur, y_axis_title = "Turnout (%)")
ggsave(file.path(root, "output/sdid_levels_turnout_figure.pdf"),
  plot = fig_turnout, height = 8, width = 10)


##
# SDID Event Study Plots
##

# Set up an event study plot function
make_event_study_plot = function(est_list, y_axis_title="XX"){
  data.frame(year = seq(1992, 2020, by=4),
    y1 = get_counterfactual(est_list$ests[["Synthetic DID (SDID)"]])$obs,
    y0_sdid = get_counterfactual(est_list$ests[["Synthetic DID (SDID)"]])$syn) %>%
    pivot_longer(
      cols = starts_with("y0_"),
      names_to = "estimator",
      names_prefix = "y0_",
      values_to = "y0"
    ) %>%
    mutate(diff = y1 - y0) %>%
    ggplot(aes(x=year, y=diff)) +
      geom_vline(xintercept = 2018, linetype = 'dotted', linewidth=1) +
      geom_line(linewidth=1.5, color="#084594") +
      geom_point(size=2, color="#084594") +
      ylab(y_axis_title) + xlab("Year") +
      ylim(-1, 1) +
      scale_x_continuous(limits = c(1991, 2021),
        breaks = seq(1992, 2020, by=4)) +
      theme(text = element_text(size=20, face="bold", color="black"),
        legend.position="bottom", legend.title= element_blank(),
        panel.grid.minor = element_blank()) +
      theme(panel.grid.minor = element_blank()) +
      scale_color_brewer(palette = 1, direction = - 1) %>%
    return()
}

# Make the vote share event study plot
event_study_vs = make_event_study_plot(allsynth_dem,
  y_axis_title = "Observed - Synthetic Dem Vote Share (%)")
ggsave(file.path(root, "output/sdid_only_event_study_demvs_figure.pdf"),
  plot = event_study_vs,
  height = 8, width = 10)

# Make the turnout event study plot
event_study_turnout = make_event_study_plot(allsynth_tur,
  y_axis_title = "Observed - Synthetic Turnout (%)")
ggsave(file.path(root, "output/sdid_only_event_study_turnout_figure.pdf"),
  plot = event_study_turnout,
  height = 8, width = 10)

